from enum import Enum
from typing import Optional, List, Dict, Any, Union
from pydantic import BaseModel, Field
from app.schemas.metrics_schema import Chart, BarChart, ScatterPlot, LineChart
from app.schemas.assets_schema import Text, Ontology, KnowledgeGraph, GraphPattern, GraphRule, MediaFile, TextRule


# 1. 任务状态
class TaskStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"


# 2. 产品类型
class ProductType(str, Enum):
    TEXT = "TEXT"
    FILE = "FILE"
    TEXT_RULE = "TEXT_RULE"
    GRAPH_RULE = "GRAPH_RULE"
    ONTOLOGY = "ONTOLOGY"
    KNOWLEDGE_GRAPH = "KNOWLEDGE_GRAPH"


# 图结构数据模型 - 节点
class Node(BaseModel):
    id: Any
    type: str
    attributes: Dict[str, Any] = Field(default_factory=dict)


# 图结构数据模型 - 边
class Edge(BaseModel):
    source: Any
    target: Any
    type: Optional[str] = "default"
    weight: Optional[float] = 1.0
    attributes: Dict[str, Any] = Field(default_factory=dict)


# 图结构数据模型
class Graph(BaseModel):
    nodes: List[Node]
    edges: List[Edge]


# 社区结构模型
class Community(BaseModel):
    id: int
    nodes: List[Any]
    type: Optional[str] = None


# 社区间演化事件
class EvolutionEvent(BaseModel):
    event_type: str  # continue, merge, split, new, dissolve
    time: int
    source_communities: Optional[List[int]] = None
    target_communities: Optional[List[int]] = None


# 时间步结果
class DHCDSTimeStepResult(BaseModel):
    time: int
    communities: Dict[Any, int]  # 节点ID到社区ID的映射
    type_communities: Dict[str, Dict[Any, int]]  # 每种节点类型的社区分配
    evolution: Dict[str, Any]  # 社区演化信息
    community_embeddings: Dict[int, List[float]]  # 社区嵌入


# DHCDS算法的输入参数
class DHCDSInputParams(BaseModel):
    graph_sequence: List[Graph]  # 多个时间步的图序列
    stability_weight: float = 0.5  # 控制社区稳定性的权重参数
    embedding_dim: int = 64  # 嵌入维度
    num_clusters: int = 10  # 默认社区数量
    node_type_weights: Optional[Dict[str, float]] = None  # 不同节点类型的权重


# DHCDS算法的输出参数
class DHCDSOutputParams(BaseModel):
    results: List[DHCDSTimeStepResult]  # 每个时间步的结果
    algorithm: str = "DHCDS"  # 算法名称
    parameters: Dict[str, Any]  # 算法参数
    community_evolution: List[Dict[str, Any]]  # 社区演化历史


# 3. 输入参数
class InputParams(BaseModel):
    dhcds_params: Optional[DHCDSInputParams] = None


# 4. 输出参数
class OutputParams(BaseModel):
    dhcds_results: Optional[DHCDSOutputParams] = None

# 5. 算法请求
class AlgorithmRequest(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: InputParams


# 6. 算法中间响应
class AlgorithmMiddleResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None

    input_params: InputParams
    error_message: Optional[str] = None

    metrics: List[Chart] = []


# 7. 算法最终响应
class AlgorithmResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None

    error_message: Optional[str] = None

    output_params: OutputParams

    metrics: List[Chart] = [] 